/**
 * @file KDTree.hpp
 * @brief 
 * @author Linfu Wei (ghowoght@qq.com)
 * @version 1.0
 * @date 2021-02-24
 * 
 * @copyright Copyright (c) 2021  WHU-EIS
 * 
 */
#ifndef _KDTREE_H
#define _KDTREE_H

#include "KDNode.hpp"
#include <iostream>
#include <vector>
#include <algorithm>
#include <memory>
#include <cmath>

using std::vector;

using KDNodePtr = std::shared_ptr<KDNode>;
using point_t = vector<double>;

class KDTree {
public:
    vector<KDNodePtr> kdnodes;
    KDNodePtr root;
    int dim = 0;

    KDTree() = default;
    KDTree(vector<point_t>& dataArray)
    {
        dim = dataArray[0].size();
        root = make_tree(dataArray.begin(), dataArray.end(), dataArray.size(), 0);
    }
    KDTree(vector<KDNodePtr>& kdnodes)
    {
        dim = kdnodes[0]->point.size();
        root = make_tree(kdnodes.begin(), kdnodes.end(), kdnodes.size(), 0);
    }

    // 先根遍历
    static void preOrderTraveral(const KDNodePtr& node)
    {
        if (node == nullptr)
            return;
        preOrderTraveral(node->left);
        preOrderTraveral(node->right);
    }

private:
    KDNodePtr make_tree(const vector<point_t>::iterator& begin,
                        const vector<point_t>::iterator& end,
                        int length,
                        int split)
    {
        if (begin == end)
            return nullptr;
        
        // => lambda表达式，实现对指定列的排序(默认升序)
        std::sort(begin, end, [split](vector<double> a, vector<double> b) {return a[split] < b[split]; });
        int middle = length / 2; // 中位数

        KDNodePtr node = std::make_shared<KDNode>(  *(begin + middle),
                                                    split,
                                                    make_tree(begin, begin + middle, middle - 1, (split + 1) % dim),
                                                    make_tree(begin + middle + 1, end, length - middle - 1, (split + 1) % dim));
        kdnodes.push_back(node);
        return node;
    }

private:
    KDNodePtr make_tree(const vector<KDNodePtr>::iterator& begin,
                        const vector<KDNodePtr>::iterator& end,
                        int length,
                        int split)
    {
        if (begin == end)
            return nullptr;
        // => lambda表达式，实现对指定列的排序(默认升序)
        std::sort(begin, end, [split](KDNodePtr a, KDNodePtr b) {return a->point[split] < b->point[split]; });
        int middle = length / 2; // 中位数
        KDNodePtr node = *(begin + middle);
        node->split = split;
        node->left = make_tree(begin, begin + middle, middle - 1, (split + 1) % dim);
        node->right = make_tree(begin + middle + 1, end, length - middle - 1, (split + 1) % dim);
        kdnodes.push_back(node);
        return node;
    }
public:
    void Add(const vector<double>& point)
    {
        KDNodePtr nodeToAdd = std::make_shared<KDNode>(point);// new KDNode(&point);
        kdnodes.push_back(nodeToAdd);
        if (root == nullptr)
            root = nodeToAdd;
        else
        {
            KDNodePtr near = GetNearest(point);
            while (true)
            {
                int split = near->split;
                if (point[split] <= near->point[split])
                {
                    if (near->left == nullptr)
                    {
                        near->left = nodeToAdd;
                        break;
                    }
                    else
                    {
                        near = near->left;
                    }
                }
                else
                {
                    if (near->right == nullptr)
                    {
                        near->right = nodeToAdd;
                        break;
                    }
                    else
                    {
                        near = near->right;
                    }
                }
            }

        }
    }

    void Add(const KDNodePtr& nodeToAdd)
    {
        kdnodes.push_back(nodeToAdd);
        if (root == nullptr)
        {
            root = nodeToAdd;

        }
        else
        {

            KDNodePtr near = GetNearest(nodeToAdd->point);
            while (true)
            {
                int split = near->split;
                if (nodeToAdd->point[split] <= near->point[split])
                {
                    if (near->left == nullptr)
                    {
                        near->left = nodeToAdd;
                        break;
                    }
                    else
                    {
                        near = near->left;
                    }
                }
                else
                {
                    if (near->right == nullptr)
                    {
                        near->right = nodeToAdd;
                        break;
                    }
                    else
                    {
                        near = near->right;
                    }
                }
            }
        }
    }
    /// 计算两个点的欧氏距离
    inline double GetDistance(const vector<double>& p1, const vector<double>& p2)
    {
        double dist = 0;
        if (p1.size() != p2.size())
            return dist;
        for (int i = 0; i < p1.size(); i++)
            dist += pow(p1[i] - p2[i], 2);
        return sqrt(dist);
    }
    inline double GetDistance(const KDNodePtr& p1, const KDNodePtr& p2)
    {
        double dist = 0;
        if (p1->dim != p2->dim)
            return dist;
        for (int i = 0; i < p1->dim; i++)
            dist += pow(p1->point[i] - p2->point[i], 2);
        return sqrt(dist);
    }

    /**
     * @brief 获取最近邻
     */
    KDNodePtr GetNearest(const vector<double>& point)
    {
        KDNodePtr best = root;
        double best_dist = GetDistance(point, best->point);
        KDNodePtr next = root;
        vector<KDNodePtr> nodes;
        // 二分查找，直到找到一个叶子节点，同时记录查找过程中的最近邻
        while (next != nullptr)
        {
            nodes.push_back(next);
            double next_dist = GetDistance(point, next->point);
            if (next_dist < best_dist)
            {
                best = next;
                best_dist = next_dist;
            }

            if (point[next->split] <= (next->point)[next->split])
                next = next->left;
            else
                next = next->right;
        }
        // 回溯，依次判断与之前节点的分割轴是否有相交
        while (nodes.size() != 0)
        {
            next = *(nodes.end() - 1);;
            nodes.pop_back();
            if (abs(point[next->split] - (next->point)[next->split]) < best_dist) // 以查询点为球心的超球面与分割超平面相交
            {
                // 进入另一个分支
                if (point[next->split] <= (next->point)[next->split])
                    next = next->right;
                else
                    next = next->left;
                // 将另一个的分支所有节点加入nodes
                while (next != nullptr)
                {
                    nodes.push_back(next);
                    double next_dist = GetDistance(point, next->point);
                    if (next_dist < best_dist)
                    {
                        best = next;
                        best_dist = next_dist;
                    }

                    if (point[next->split] <= (next->point)[next->split])
                        next = next->left;
                    else
                        next = next->right;
                }
            }
        }
        return best;
    }

    /**
     * @brief 获取输入节点周围指定范围内的节点
     */
    void GetRange(vector<KDNodePtr>& neighbors, point_t& point, double R)
    {
        KDNodePtr next = root;
        vector<KDNodePtr> nodes;
        // 二分查找
        while (next != nullptr)
        {
            nodes.push_back(next);
            double next_dist = GetDistance(point, next->point);
            if (next_dist < R)
            {
                neighbors.push_back(next);
            }

            if (point[next->split] <= next->point[next->split])
                next = next->left;
            else
                next = next->right;
        }
        // 回溯
        while (nodes.size() != 0)
        {
            next = *(nodes.end() - 1);
            nodes.pop_back();
            if (abs(point[next->split] - next->point[next->split]) < R) // 以查询点为球心、R为半径的超球面与分割超平面相交
            {
                // 进入另一个分支
                if (point[next->split] <= next->point[next->split])
                    next = next->right;
                else
                    next = next->left;
                // 将另一个的分支所有节点加入nodes
                while (next != nullptr)
                {
                    nodes.push_back(next);
                    double next_dist = GetDistance(point, next->point);
                    if (next_dist < R)
                        neighbors.push_back(next);

                    if (point[next->split] <= next->point[next->split])
                        next = next->left;
                    else
                        next = next->right;
                }
            }
        }
    }

};


#endif